{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Jxv6goXm7oGF"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "llMNufAK7nfK"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8Byow2J6LaPl"
      },
      "source": [
        "# tf.function and AutoGraph in TensorFlow 2.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kGXS3UWBBNoc"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/2/guide/autograph\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/2/guide/autograph.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/2/guide/autograph.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CydFK2CL7ZHA"
      },
      "source": [
        "TF 2.0 brings together the ease of eager execution and the power of TF 1.0. At the center of this merger is `tf.function`, which allows you to transform a subset of Python syntax into portable, high-performance TensorFlow graphs.\n",
        "\n",
        "A cool new feature of `tf.function` is AutoGraph, which lets you write graph code using natural Python syntax. For a list of the Python features that you can use with AutoGraph, see [AutoGraph Capabilities and Limitations](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/LIMITATIONS.md). For more details about `tf.function`, see the RFC [TF 2.0: Functions, not Sessions](https://github.com/tensorflow/community/blob/master/rfcs/20180918-functions-not-sessions-20.md). For more details about AutoGraph, see `tf.autograph`.\n",
        "\n",
        "This tutorial will walk you through the basic features of `tf.function` and AutoGraph."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n4EKOpw9mObL"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Import TensorFlow 2.0 Preview Nightly and enable TF 2.0 mode:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "V9oECvVSI1Kj"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mT7meGqrZTz9"
      },
      "outputs": [],
      "source": [
        "!pip install tf-nightly-2.0-preview\n",
        "import tensorflow as tf"
      ]
    },
    {
      "metadata": {
        "id": "ypIdFDCTfjLj",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Install a temporary patch to enable a few extra TF 2.0 upgrades. This piece will be removed soon."
      ]
    },
    {
      "metadata": {
        "id": "IfgTXZETfQ4u",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from tensorflow.python.ops import control_flow_util\n",
        "control_flow_util.ENABLE_CONTROL_FLOW_V2 = True"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "77AsVr1GGtBP"
      },
      "source": [
        "## The `tf.function` decorator\n",
        "\n",
        "When you annotate a function with `tf.function`, you can still call it like any other function. But it will be compiled into a graph, which means you get the benefits of faster execution, running on GPU or TPU, or exporting to SavedModel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "FhIg7-z6HNWj"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def simple_nn_layer(x, y):\n",
        "  return tf.nn.relu(tf.matmul(x, y))\n",
        "\n",
        "\n",
        "x = tf.random.uniform((3, 3))\n",
        "y = tf.random.uniform((3, 3))\n",
        "\n",
        "simple_nn_layer(x, y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "U-LAE4pMNR9g"
      },
      "source": [
        "If we examine the result of the annotation, we can see that it's a special callable that handles all interactions with the TensorFlow runtime. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "q4t2iuS7Nqc0"
      },
      "outputs": [],
      "source": [
        "simple_nn_layer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DqeefLGNXjZQ"
      },
      "source": [
        "If your code uses multiple functions, you don't need to annotate them all - any functions called from an annotated function will also run in graph mode."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3VGF7tlVXiZY"
      },
      "outputs": [],
      "source": [
        "def linear_layer(x):\n",
        "  return 2 * x + 1\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def deep_net(x):\n",
        "  return tf.nn.relu(linear_layer(x))\n",
        "\n",
        "\n",
        "deep_net(tf.constant((1, 2, 3)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ohbSnA79mcJV"
      },
      "source": [
        "## Use Python control flow\n",
        "\n",
        "When using data-dependent control flow inside `tf.function`, you can use Python control flow statements and AutoGraph will convert them into appropriate TensorFlow ops. For example, `if` statements will be converted into `tf.cond()` if they depend on a `Tensor`.\n",
        "\n",
        "In the example below, `x` is a `Tensor` but the `if` statement works as expected:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aA3gOodCBkOw"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def square_if_positive(x):\n",
        "  if x \u003e 0:\n",
        "    x = x * x\n",
        "  else:\n",
        "    x = 0\n",
        "  return x\n",
        "\n",
        "\n",
        "print('square_if_positive(2) = {}'.format(square_if_positive(tf.constant(2))))\n",
        "print('square_if_positive(-2) = {}'.format(square_if_positive(tf.constant(-2))))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "GMiCUkdyoq98"
      },
      "source": [
        "Note: the example above shows how to perform simple conditionals when scalar values are involves. Typical ML code involves batches; in those cases you should consider using the faster and vectorized `tf.where` if possible."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "m-jWmsCmByyw"
      },
      "source": [
        "AutoGraph supports common Python statements like `while`, `for`, `if`, `break`, `continue` and `return`, with support for nesting. That means you can use `Tensor` expressions in the condition of `while` and `if` statements, or iterate over a `Tensor` in a `for` loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "toxKBOXbB1ro"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def sum_even(items):\n",
        "  s = 0\n",
        "  for c in items:\n",
        "    if c % 2 \u003e 0:\n",
        "      continue\n",
        "    s += c\n",
        "  return s\n",
        "\n",
        "\n",
        "sum_even(tf.constant([10, 12, 15, 20]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "AtDaLrbySw4j"
      },
      "source": [
        "AutoGraph also provides a low-level API for advanced users. For example we can use it to have a look at the generated code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aRsde3x_SjTQ"
      },
      "outputs": [],
      "source": [
        "print(tf.autograph.to_code(sum_even.python_function, experimental_optional_features=None))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rvJXCfk8VkLf"
      },
      "source": [
        "Here's an example of more complicated control flow:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "h-Z87IJqVlKl"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def fizzbuzz(n):\n",
        "  msg = tf.constant('')\n",
        "  for i in range(n):\n",
        "    if i % 3 == 0:\n",
        "      msg += 'Fizz'\n",
        "    elif i % 5 == 0:\n",
        "      msg += 'Buzz'\n",
        "    else:\n",
        "      msg += tf.as_string(i)\n",
        "    msg += '\\n'\n",
        "  return msg\n",
        "\n",
        "\n",
        "print(fizzbuzz(tf.constant(15)).numpy().decode())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-pkEH6OecW7h"
      },
      "source": [
        "## Use Python `print`\n",
        "\n",
        "AutoGraph will also convert Python builtins like `print`.\n",
        "\n",
        "Note: due to the parallel nature of calculations in TensorFlow, statements might execute out of order. It's best to use `print` only to inspect actual values, and you should not use it to determine whether the program execution reaches a certain point."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ehBac9rUR6nh"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def count(n):\n",
        "  for i in tf.range(n):\n",
        "    print(i)\n",
        "\n",
        "\n",
        "count(tf.constant(5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "OlWcPngbBZ4-"
      },
      "source": [
        "## Other handy conversions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "cf4hLjRaABTJ"
      },
      "source": [
        "Other builtins that AutoGraph can adapt for TensorFlow`range` and `len`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "a60k91mvBFsu"
      },
      "source": [
        "`range` is a shortcut for `tf.range`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_Lm9ugqUAORh"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def range_example(n):\n",
        "  return range(n)\n",
        "\n",
        "\n",
        "print(range_example(tf.constant(3)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5ZdA2Ze9BJ1n"
      },
      "source": [
        "`len` is a shortcut for `.shape[0]`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-7UD-BJZAxyW"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def len_example(n):\n",
        "  return len(n)\n",
        "\n",
        "\n",
        "print(len_example(tf.zeros((20, 10))))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "h_Y4uC1R1B55"
      },
      "source": [
        "## Keras and AutoGraph\n",
        "\n",
        "You can use `tf.function` with object methods as well. For example, you can decorate your custom Keras models, typically by annotating the model's `call` function. For more information, see `tf.keras`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cR6mpLKP1HLe"
      },
      "outputs": [],
      "source": [
        "class CustomModel(tf.keras.models.Model):\n",
        "  \n",
        "  @tf.function\n",
        "  def call(self, input_data):\n",
        "    if tf.reduce_mean(input_data) \u003e 0:\n",
        "      return input_data\n",
        "    else:\n",
        "      return input_data // 2\n",
        "\n",
        "\n",
        "model = CustomModel()\n",
        "\n",
        "model(tf.constant([-2, -4]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NTEvpBK9f8kj"
      },
      "source": [
        "## Side effects\n",
        "\n",
        "Just like in eager mode, you can use operations with side effects, like `tf.assign` or `tf.print` normally inside `tf.function`, and it will insert the necessary control dependencies to ensure they execute in order."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-Wd6i8S9gcuC"
      },
      "outputs": [],
      "source": [
        "v = tf.Variable(5)\n",
        "\n",
        "@tf.function\n",
        "def find_next_odd():\n",
        "  v.assign(v + 1)\n",
        "  if v % 2 == 0:\n",
        "    v.assign(v + 1)\n",
        "\n",
        "\n",
        "find_next_odd()\n",
        "v"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4LfnJjm0Bm0B"
      },
      "source": [
        "## Example: training a simple model\n",
        "\n",
        "AutoGraph also allows you to move more computation inside TensorFlow. For example, a training loop is just control flow, so it can actually be brought into TensorFlow."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Em5dzSUOtLRP"
      },
      "source": [
        "### Download data"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "xqoxumv0ssQW",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def prepare_mnist_features_and_labels(x, y):\n",
        "  x = tf.cast(x, tf.float32) / 255.0\n",
        "  y = tf.cast(y, tf.int64)\n",
        "  return x, y\n",
        "\n",
        "def mnist_dataset():\n",
        "  (x, y), _ = tf.keras.datasets.mnist.load_data()\n",
        "  ds = tf.data.Dataset.from_tensor_slices((x, y))\n",
        "  ds = ds.map(prepare_mnist_features_and_labels)\n",
        "  ds = ds.take(20000).shuffle(20000).batch(100)\n",
        "  return ds\n",
        "\n",
        "train_dataset = mnist_dataset()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "znmy4l8ntMvW"
      },
      "source": [
        "### Define the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ltxyJVWTqNAO"
      },
      "outputs": [],
      "source": [
        "model = tf.keras.Sequential((\n",
        "    tf.keras.layers.Reshape(target_shape=(28 * 28,), input_shape=(28, 28)),\n",
        "    tf.keras.layers.Dense(100, activation='relu'),\n",
        "    tf.keras.layers.Dense(100, activation='relu'),\n",
        "    tf.keras.layers.Dense(10)))\n",
        "model.build()\n",
        "optimizer = tf.keras.optimizers.Adam()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "oeYV6mKnJGMr"
      },
      "source": [
        "### Define the training loop"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "3xtg_MMhJETd",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def compute_loss(logits, labels):\n",
        "  return tf.reduce_mean(\n",
        "      tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
        "          logits=logits, labels=labels))\n",
        "\n",
        "\n",
        "def compute_accuracy(logits, labels):\n",
        "  predictions = tf.argmax(logits, axis=1)\n",
        "  return tf.reduce_mean(tf.cast(predictions == labels, tf.float32))\n",
        "\n",
        "\n",
        "def train_one_step(model, optimizer, x, y):\n",
        "  with tf.GradientTape() as tape:\n",
        "    tape.watch(model.variables)\n",
        "    logits = model(x)\n",
        "    loss = compute_loss(logits, y)\n",
        "\n",
        "  grads = tape.gradient(loss, model.variables)\n",
        "  optimizer.apply_gradients(zip(grads, model.variables))\n",
        "\n",
        "  accuracy = compute_accuracy(logits, y)\n",
        "  return loss, accuracy\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def train(model, optimizer, train_ds):\n",
        "  step = 0\n",
        "  for x, y in train_ds:\n",
        "    step += 1\n",
        "    loss, accuracy = train_one_step(model, optimizer, x, y)\n",
        "    if step % 10 == 0:\n",
        "      print('Step', step, ': loss', loss, '; accuracy', accuracy)\n",
        "  return step\n",
        "\n",
        "_ = train(model, optimizer, train_dataset)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SnsumiP6eRYL"
      },
      "source": [
        "## A note on batching\n",
        "\n",
        "In real applications batching is essential for performance. The best code to convert to AutoGraph is code where the control flow is decided at the _batch_ level. If making decisions at the individual _example_ level, try to use batch APIs to maintain performance.\n",
        "\n",
        "For example, if you have the following code in Python:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "t31QoERiNccJ"
      },
      "outputs": [],
      "source": [
        "def square_if_positive(x):\n",
        "  return [i ** 2 if i \u003e 0 else i for i in x]\n",
        "\n",
        "\n",
        "square_if_positive(range(-5, 5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kSeEJ76uNgwD"
      },
      "source": [
        "You may be tempted to write it in TensorFlow as such (and this would work!):\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RqR8WzSzNf87"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def square_if_positive_naive(x):\n",
        "  result = tf.TensorArray(tf.int32, size=len(x))\n",
        "  for i in range(len(x)):\n",
        "    if x[i] \u003e 0:\n",
        "      result = result.write(i, x[i] ** 2)\n",
        "    else:\n",
        "      result = result.write(i, x[i])\n",
        "  return result.stack()\n",
        "\n",
        "\n",
        "square_if_positive_naive(tf.range(-5, 5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gTcyWXVGN3gS"
      },
      "source": [
        "But in this case, it turns out you can write the following:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VO2f6x-lNfVj"
      },
      "outputs": [],
      "source": [
        "def square_if_positive_vectorized(x):\n",
        "  return tf.where(x \u003e 0, x ** 2, x)\n",
        "\n",
        "\n",
        "square_if_positive_vectorized(tf.range(-5, 5))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "Jxv6goXm7oGF"
      ],
      "last_runtime": {
        "build_target": "",
        "kind": "local"
      },
      "name": "TF2.0 Functions and Autograph",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1Hl4PR32Y_nhq4AZI-ku9FlpZIgXjEU7d",
          "timestamp": 1544141357531
        },
        {
          "file_id": "https://github.com/tensorflow/docs/blob/master/site/en/guide/autograph.ipynb",
          "timestamp": 1542126584672
        }
      ],
      "toc_visible": true,
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
